// Copyright 2022 Chen Jun
// 在MIT许可证下发布。

#ifndef ARMOR_DETECTOR_ARMOR_HPP_
#define ARMOR_DETECTOR_ARMOR_HPP_

#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
#include <Eigen/Eigen>

// STL
#include <algorithm>
#include <string>

namespace hnurm
{
    const int BLUE = 0;
    const int RED = 1;

    enum ArmorType
    {
        SMALL = 0, LARGE = 1
    };

    /**
     * @brief Light类表示一个灯，继承自cv::RotatedRect类
     */
    struct Light: public cv::RotatedRect
    {
        Light() = default;

        /**
         * @brief 构造函数
         * @param box 矩形框
         */
        explicit Light(const cv::RotatedRect& box)
            : cv::RotatedRect(box)
        {
            cv::Point2f p[4];
            box.points(p);
            std::sort(p, p + 4, [](const cv::Point2f &a, const cv::Point2f &b)
                      { return a.y < b.y; });
            top = (p[0] + p[1]) / 2;
            bottom = (p[2] + p[3]) / 2;

            length = cv::norm(top - bottom);
            width = cv::norm(p[0] - p[1]);

            tilt_angle = std::atan2(std::abs(top.x - bottom.x), std::abs(top.y - bottom.y));
            tilt_angle = tilt_angle / CV_PI * 180;
        }

        int color;                    // 灯的颜色
        cv::Point2f top, bottom;       // 灯的顶点和底点
        double length;                 // 灯的长度
        double width;                  // 灯的宽度
        float tilt_angle;              // 灯的倾斜角度
    };

    /**
     * @brief Armor类表示一个装甲，包含灯、中心点、四个角点、数字ROI图像、数字、索引、相似度、置信度、分类结果、装甲类型等信息
     */
    struct Armor
    {
        // 左灯条、右灯条
        Light left_light, right_light;

        // 装甲板中心点
        cv::Point2f center;

        // 装甲板的四个角点，顺序为左上、左下、右下、右上
        std::vector<cv::Point2f> points2d;

        // 数字ROI图像
        cv::Mat number_img;

        std::string number;          // 数字
        int idx{};                   // 索引
        float similarity{};          // 相似度
        float confidence{};          // 置信度
        std::string classification_result;   // 分类结果

        ArmorType armor_type = ArmorType::SMALL;  // 装甲类型

        double distance_to_image_center{};      // 距离图像中心的像素值
        // todo
        Eigen::Vector3d position;           // 位置
        Eigen::Matrix3d rotation;            // 旋转
        Eigen::Matrix3d _r1;                 // 相机坐标系到云台坐标系
        Eigen::Matrix3d _r2;                 // 先对pitch进行旋转
        Eigen::Matrix3d _r3;                 // 再对yaw进行旋转，变为世界坐标系（pitch、yaw为0）
        Eigen::Matrix3d _rmat;
        Eigen::Vector3d _translation;

        Armor() = default;

        /**
         * @brief 构造函数
         * @param l1 左灯条
         * @param l2 右灯条
         */
        Armor(const Light &l1, const Light &l2)
        {
            if (l1.center.x < l2.center.x)
            {
                left_light = l1, right_light = l2;
            }
            else
            {
                left_light = l2, right_light = l1;
            }
            center = (left_light.center + right_light.center) / 2;

            points2d.resize(9);
            points2d[0] = left_light.top;
            points2d[1] = left_light.bottom;
            points2d[2] = right_light.bottom;
            points2d[3] = right_light.top;

            points2d[4] = (left_light.top+right_light.top)/2;
            points2d[5] = (left_light.top+left_light.bottom)/2;
            points2d[6] = (left_light.bottom+right_light.bottom)/2;
            points2d[7] = (right_light.bottom+right_light.top)/2;

            points2d[8] = center;
        }

        /**
         * @brief 在给定的图像上绘制装甲
         * @param canvas 绘制的图像
         */
        void showArmor(cv::Mat &canvas)
        {
            // 绘制装甲
            cv::line(canvas, points2d[0], points2d[1], cv::Scalar(0, 255, 0), 2);
            cv::line(canvas, points2d[1], points2d[2], cv::Scalar(0, 255, 0), 2);
            cv::line(canvas, points2d[2], points2d[3], cv::Scalar(0, 255, 0), 2);
            cv::line(canvas, points2d[3], points2d[0], cv::Scalar(0, 255, 0), 2);
            // 绘制中心点
            cv::circle(canvas, center, 10, cv::Scalar(0, 255, 0), 2);
            // 在装甲的左上角绘制分类结果
            cv::putText(canvas, classification_result, points2d[0], cv::FONT_HERSHEY_SIMPLEX,
                        1, cv::Scalar(0, 255, 0), 2);

        }
    };

}  // namespace hnurm

#endif  // ARMOR_DETECTOR_ARMOR_HPP_
